from torch import nn
from torch import optim
from torch.optim import lr_scheduler

import ai

#  hello
if __name__ == "__main__":
    model = ai.store_model.store_model.load_model(ai.store_model.store_model_name.TEXTCLASSIFY_1)
    optim = optim.Adam(model.parameters(), lr=0.001)
    loss_f = nn.CrossEntropyLoss()
    lr_schedule = lr_scheduler.CosineAnnealingLR(optim, T_max=1, eta_min=1e-10, last_epoch=-1)
    train_loader, valid_loader, test_loader = ai.store_data.store_data.DataStore(
        ai.store_data.store_data_name.TEXT_CLASSIFY).get_dataloader()
    runner = ai.run.runner.RunnerGxl(model, optim, loss_f, valid_loader, test_loader, lr_schedule,
                                     model_name=ai.store_model.store_model_name.TEXTCLASSIFY_1)
    runner.run(10)
